import tqdm
from multiprocessing import Queue, Condition, Process
from typing import Callable
from threading import Thread
from backend.experiment.framework.log import Log


class ListProgressVisibleMultiprocessing:
    """
    生产者消费者模型下的带进度条的多进程任务执行器
    """

    def __init__(self, processes: int, process_func: Callable):
        self.__processes = processes
        self.__process_func = process_func

    def process(
            self, todo_list: list, batch, verbose: bool = True,
            kwargs: dict = None
    ):
        kwargs = {} if kwargs is None else kwargs
        has_task = Condition()
        no_task = Condition()
        task_queue = Queue()
        result_queue = Queue()

        processes = []
        for i in range(self.__processes):
            target_kwargs = kwargs.copy()
            target_kwargs.update({
                'process_id': i,
                'task_queue': task_queue,
                'has_task': has_task,
                'no_task': no_task,
                'result_queue': result_queue,
            })
            processes.append(
                Process(target=self._process_main, kwargs=target_kwargs)
            )

        for each in processes:
            each.start()

        task_send_thread = Thread(target=self._send_task_thread_main, kwargs={
            'task_list': todo_list,
            'task_queue': task_queue,
            'has_task': has_task,
            'no_task': no_task,
            'batch': batch
        })
        task_send_thread.start()

        result_list_dict = {}
        with tqdm.tqdm(range(len(todo_list)), disable=not verbose) as bar:
            while len(result_list_dict) < len(todo_list):
                completed = result_queue.get()
                for index, content in completed:
                    result_list_dict[index] = content
                    bar.update(1)

        Log(f'main-process waiting worker join', -1)
        for index, each in enumerate(processes):
            each.join()
            Log(f'process-{index} joined', -1)
        task_send_thread.join()
        result = []
        for i in range(len(todo_list)):
            result.append(result_list_dict[i])
        return result

    def _send_task_thread_main(
            self, task_list: list,
            task_queue: Queue,
            has_task: Condition,
            no_task: Condition,
            batch: int
    ):
        done = False
        begin = 0
        put_tasks_each = self.__processes + 1
        end = min(put_tasks_each * batch, len(task_list))
        while not done:
            putting_batches = []
            batched_tasks = []
            i = begin
            while len(putting_batches) < put_tasks_each and i < end:
                if len(batched_tasks) >= batch:
                    putting_batches.append(batched_tasks)
                    batched_tasks = []
                batched_tasks.append((i, task_list[i]))
                i += 1
            if len(batched_tasks) > 0:
                putting_batches.append(batched_tasks)
            Log(f'main-send-thread wait task queue empty', -1)
            with no_task:
                while task_queue.qsize() >= self.__processes + 1:
                    no_task.wait()
                Log(f'main-send-thread putting [{begin}:{end}]', -1)
                for each in putting_batches:
                    task_queue.put(each)

                if end == len(task_list):
                    done = True
                    task_queue.put([])
                    Log(f'main-send-thread task set finish signal', -1)

            with has_task:
                Log(f'main-send-thread notify work processes', -1)
                has_task.notify_all()
            begin = end
            end = min(end + put_tasks_each * batch, len(task_list))
        Log(f'main-send-thread exit', -1)

    def _process_main(
            self,
            process_id: int,
            task_queue: Queue,
            has_task: Condition,
            no_task: Condition,
            result_queue: Queue,
            **kwargs
    ):
        while True:
            Log(f'process-{process_id} trying getting tasks,'
                f' batch remain: {task_queue.qsize()}', -1)
            with has_task:
                while task_queue.qsize() <= 0:
                    with no_task:
                        no_task.notify_all()
                    has_task.wait()
                batch_inputs = task_queue.get()
            if len(batch_inputs) == 0:
                Log(f'process-{process_id} receive finished signal', -1)
                task_queue.put([])
                with has_task:
                    has_task.notify_all()
                break
            result_list = []
            Log(f'process-{process_id} handling batch task,'
                f' len: {len(batch_inputs)}', -1)
            for index, target in batch_inputs:
                result_list.append(
                    (index, self.__process_func(target, **kwargs))
                )
            if len(result_list) > 0:
                result_queue.put(result_list)
        Log(f'process-{process_id} exit', -1)
